from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass
from typing import Dict


@dataclass
class TokenBucket:
    rate_per_min: int
    capacity: int
    tokens: float
    last_refill: float

    def consume(self, amount: int = 1) -> bool:
        now = time.monotonic()
        elapsed = now - self.last_refill
        if elapsed > 0:
            refill_amount = (self.rate_per_min / 60.0) * elapsed
            self.tokens = min(self.capacity, self.tokens + refill_amount)
            self.last_refill = now
        if self.tokens >= amount:
            self.tokens -= amount
            return True
        return False


class RateLimiter:
    def __init__(self) -> None:
        self._buckets: Dict[str, TokenBucket] = {}
        self._lock = asyncio.Lock()

    async def allow(self, key: str, rate_per_min: int) -> bool:
        if rate_per_min <= 0:
            return True
        async with self._lock:
            bucket = self._buckets.get(key)
            if bucket is None:
                bucket = TokenBucket(
                    rate_per_min=rate_per_min,
                    capacity=rate_per_min,
                    tokens=float(rate_per_min),
                    last_refill=time.monotonic(),
                )
                self._buckets[key] = bucket
            return bucket.consume()


script_rate_limiter = RateLimiter()
user_rate_limiter = RateLimiter()

